import pandas as pd
from _data_process import cumulative_to_monthly
from _data_process import monthly_to_quarterly
from _data_process import calculate_yoy

from _data_plot import plot_double_axis_with_background


def merge_series_with_common_dates(series1, series2, col_name1, col_name2):
    """
    将两个Series合并为DataFrame，保留共同日期及最新两期非共同日期的数据
    
    参数:
    series1: 第一个Series，索引为日期
    series2: 第二个Series，索引为日期
    col_name1: 第一个Series在DataFrame中的列名
    col_name2: 第二个Series在DataFrame中的列名
    
    返回:
    DataFrame: 合并后的结果，包含两个Series的数据
    """
    # 转换为DataFrame并指定列名
    df1 = series1.to_frame(name=col_name1)
    df2 = series2.to_frame(name=col_name2)
    
    # 合并两个DataFrame，保留所有日期
    merged = df1.join(df2, how='outer')
    
    # 获取两个Series的日期索引
    index1 = set(series1.index)
    index2 = set(series2.index)
    common_dates = index1 & index2  # 共同日期
    
    # 获取每个Series的最新两期日期（按索引排序后取最后两个）
    latest1 = set(series1.sort_index().index[-2:]) if not series1.empty else set()
    latest2 = set(series2.sort_index().index[-2:]) if not series2.empty else set()
    
    # 需要保留的日期：共同日期 + 各自最新两期
    keep_dates = common_dates | latest1 | latest2
    
    # 筛选出需要保留的日期并按日期排序
    result = merged.loc[list(keep_dates)].sort_index()
    
    return result.round(2)


def add_same_trend_column(df):
    """
    为DataFrame添加四列，分别表示第一列涨跌、第二列涨跌、是否同趋势以及index的月份
    上涨为1，下跌为-1，不变为0；同时上涨或同时下跌输出1，否则输出0
    
    参数:
    df: 包含两列数据的DataFrame，前两列分别为需要比较的指标，索引为日期
    
    返回:
    DataFrame: 添加了"第一列涨跌"、"第二列涨跌"、"同趋势"和"月份"列的新DataFrame
    """
    # 计算每列的环比变化（当前值与前一期的差值）
    diff1 = df.iloc[:, 0].diff()  # 第一列的变化
    diff2 = df.iloc[:, 1].diff()  # 第二列的变化
    
    name1 = df.columns[0]
    name2 = df.columns[1]
    # 判断第一列涨跌：上涨为1，下跌为-1，不变为0
    df_with_trend = df.copy()
    df_with_trend[name1 + "涨跌"] = 0
    df_with_trend.loc[diff1 > 0, name1 + "涨跌"] = 1
    df_with_trend.loc[diff1 < 0, name1 + "涨跌"] = -1
    
    # 判断第二列涨跌：上涨为1，下跌为-1，不变为0
    name_up_down = name2 + "涨跌"
    df_with_trend[name_up_down] = 0
    df_with_trend.loc[diff2 > 0, name_up_down] = 1
    df_with_trend.loc[diff2 < 0, name_up_down] = -1
    
    # 判断是否同时为正（都上涨）或同时为负（都下跌）
    # 排除两者都为0或一正一负的情况
    same_trend = ((diff1 > 0) & (diff2 > 0)) | ((diff1 < 0) & (diff2 < 0))
    
    # 将布尔值转换为1和0
    df_with_trend['同趋势'] = same_trend.astype(int)
    
    # 新增"月份"列，提取index的月份（假设index为datetime类型）
    df_with_trend['月份'] = df.index.month
    df_with_trend['xx'] = df_with_trend[name_up_down]
    return df_with_trend

if __name__ == "__main__":
    elec_data_cum = pd.read_excel(f"data/全社会用电量当月累计值.xlsx", index_col=0)
    gdp_data = pd.read_excel(f"data/GDP_不变价_当季同比.xlsx", index_col=0)

    # -----------------------------------------------------------
    # 计算用电量的季度同比
    monthly_data = cumulative_to_monthly(elec_data_cum) # 在原始数据是累计值的时候用
    quarterly_data = monthly_to_quarterly(monthly_data)
    elec_data_quarterly_yoy = calculate_yoy(quarterly_data)

    gdp_elec_data = merge_series_with_common_dates(
        gdp_data.iloc[:,0], elec_data_quarterly_yoy.iloc[:,0], 
        "GDP当季同比", "全社会用电量当季同比").iloc[5:,:]

    gdp_elec_data = add_same_trend_column(gdp_elec_data)
    # dic = {month: df for month, df in gdp_elec_data.groupby('月份')}

    fig1 = plot_double_axis_with_background(
        gdp_elec_data.iloc[:,:], 
        gdp_elec_data.columns[0], 
        gdp_elec_data.columns[1],
        title="GDP与用电量关联分析",
        )
    # -----------------------------------------------------------

    # pmi_data = pd.read_excel(f"data/PMI.xlsx", index_col=0)
    # gdp_pmi_data = merge_series_with_common_dates(
    #     gdp_data.iloc[:,0], pmi_data.iloc[:,1], 
    #     "GDP当季同比", "PMI").iloc[5:,:]
    # gdp_pmi_data = add_same_trend_column(gdp_pmi_data)

    # fig2 = plot_double_axis_with_background(
    #     gdp_pmi_data.iloc[:,:], 
    #     gdp_pmi_data.columns[0], 
    #     gdp_pmi_data.columns[1],
    #     title="GDP与PMI关联分析",
    #     )
    # -----------------------------------------------------------

    gongye_data = pd.read_excel(f"data/工业增加值当月同比.xlsx", 
                                index_col=0)
    quarter_gongye_data = monthly_to_quarterly(gongye_data)

    gdp_gongye_data = merge_series_with_common_dates(
        gdp_data.iloc[:,0], quarter_gongye_data.iloc[:,0], 
        "GDP当季同比", "工业增加值当月同比当季均值").iloc[5:,:]
    gdp_gongye_data = add_same_trend_column(gdp_gongye_data)
    fig3 = plot_double_axis_with_background(
        gdp_gongye_data.iloc[:,:], 
        gdp_gongye_data.columns[0], 
        gdp_gongye_data.columns[1],
        title="GDP与工业增加值关联分析",
        )

    # -----------------------------------------------------------
    seveice_data = pd.read_excel(f"data/服务业生产指数当月同比.xlsx", 
                                index_col=0)
    quarter_seveice_data = monthly_to_quarterly(seveice_data)
    quarter_seveice_data = quarter_seveice_data
    gdp_service_data = merge_series_with_common_dates(
        gdp_data.iloc[:,0], quarter_seveice_data.iloc[:,0], 
        "GDP当季同比", "服务业生产指数当月同比当季均值")
    gdp_service_data = add_same_trend_column(gdp_service_data)
    fig4 = plot_double_axis_with_background(
        gdp_service_data.iloc[:,:], 
        gdp_service_data.columns[0], 
        gdp_service_data.columns[1],
        title="GDP与服务业生产指数关联分析",
        )

    # -----------------------------------------------------------
    house_data = pd.read_excel(f"data/30大中城市商品房成交面积_当周值.xlsx", 
                                index_col=0)
    quarter_house_data = monthly_to_quarterly(house_data)
    quarter_house_data_yoy = calculate_yoy(quarter_house_data)
    gdp_house_data = merge_series_with_common_dates(
        gdp_data.iloc[:,0], quarter_house_data_yoy.iloc[:,0], 
        "GDP当季同比", "30大中城市商品房成交面积当季同比")
    gdp_house_data = add_same_trend_column(gdp_house_data).iloc[5:,:]
    fig5 = plot_double_axis_with_background(
        gdp_house_data.iloc[:,:], 
        gdp_house_data.columns[0], 
        gdp_house_data.columns[1],
        title="GDP与30大中城市商品房成交面积关联分析",
        )

    # -----------------------------------------------------------
    movie_data = pd.read_excel(f"data/电影票房收入当日值.xlsx", 
                                index_col=0) # 首先准备好excel，从哪儿下都行，wind结构化更好
    quarter_movie_data = monthly_to_quarterly(movie_data) # 当日值
    quarter_movie_data_yoy = calculate_yoy(quarter_movie_data) # 计算季度同比
    gdp_movie_data = merge_series_with_common_dates(
        gdp_data.iloc[:,0], quarter_movie_data_yoy.iloc[:,0], 
        "GDP当季同比", "电影票房收入当季同比")
    gdp_movie_data = add_same_trend_column(gdp_movie_data).iloc[5:,:]
    fig_movie = plot_double_axis_with_background(
        gdp_movie_data.iloc[:,:], 
        gdp_movie_data.columns[0], 
        gdp_movie_data.columns[1],
        title="GDP与电影票房收入关联分析",
        )

    # -----------------------------------------------------------
    sh_subway_data = pd.read_excel(f"data/上海地铁客运量当日值.xlsx", 
                                index_col=0) # 首先准备好excel，从哪儿下都行，wind结构化更好
    quarter_sh_subway_data = monthly_to_quarterly(sh_subway_data) # 当日值
    quarter_sh_subway_data_yoy = calculate_yoy(quarter_sh_subway_data) # 计算季度同比   
    gdp_sh_subway_data = merge_series_with_common_dates(
        gdp_data.iloc[:,0], quarter_sh_subway_data_yoy.iloc[:,0], 
        "GDP当季同比", "上海地铁客运量当季同比")
    gdp_sh_subway_data = add_same_trend_column(gdp_sh_subway_data).iloc[5:,:]
    fig_sh_subway = plot_double_axis_with_background(
        gdp_sh_subway_data.iloc[:,:], 
        gdp_sh_subway_data.columns[0], 
        gdp_sh_subway_data.columns[1],
        title="GDP与上海地铁客运量关联分析",
        )

    # ---------------------------------------------------------
    sales_data = pd.read_excel(f"data/社会消费品零售总额当月同比.xlsx", 
                                index_col=0)
    quarter_sales_data = monthly_to_quarterly(sales_data)
    gdp_sales_data = merge_series_with_common_dates(
        gdp_data.iloc[:,0], quarter_sales_data.iloc[:,0], 
        "GDP当季同比", "社会消费品零售总额当季同比")
    gdp_sales_data = add_same_trend_column(gdp_sales_data)
    fig_sales = plot_double_axis_with_background(
        gdp_sales_data.iloc[:,:], 
        gdp_sales_data.columns[0], 
        gdp_sales_data.columns[1],
        title="GDP与社会消费品零售总额关联分析",
        )